package com.morning.aiagentbackend.chatmemory;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import org.objenesis.strategy.StdInstantiatorStrategy;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

/**
 * 基于文件存储的会话历史记录
 */
public class FileBasedChatMemory implements ChatMemory {

    private final String BASE_DIR;

    private static final Kryo KRYO = new Kryo();

    static {
        KRYO.setRegistrationRequired(false);
        // 设置实例化策略
        KRYO.setInstantiatorStrategy(new StdInstantiatorStrategy());
    }

    // 构造对象时，指定存储目录
    public FileBasedChatMemory(String baseDir) {
        this.BASE_DIR = baseDir;
        // 创建目录
        File file = new File(baseDir);
        if (!file.exists()) {
            file.mkdirs();
        }
    }


    @Override
    public void add(String conversationId, Message message) {
        this.add(conversationId, List.of(message));
    }

    @Override
    public void add(String conversationId, List<Message> messages) {
        List<Message> messageList = getOrCreateConversation(conversationId);
        messageList.addAll(messages);
        saveConversation(conversationId, messageList);
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        List<Message> messages = getOrCreateConversation(conversationId);
        return messages.stream()
                .skip(Math.max(0, messages.size() - lastN))
                .toList();
    }

    @Override
    public void clear(String conversationId) {
        File conversationFile = getConversationFile(conversationId);
        if (conversationFile.exists()) {
            conversationFile.delete();
        }
    }


    /**
     * 根据会话ID获取对应的会话文件对象
     * 该方法用于构造一个指向特定会话文件的File对象它将BASE_DIR作为父目录，
     * 并结合会话ID与".kryo"后缀来形成文件名这种方法便于统一管理和定位会话数据文件
     *
     * @param conversationId 会话的唯一标识符，用于区分不同的会话文件
     * @return File对象，表示特定会话的数据文件
     */
    private File getConversationFile(String conversationId) {
        return new File(BASE_DIR, conversationId + ".kryo");
    }

    /**
     * 获取或创建对话
     * 如果对话文件已存在，则从文件中读取对话记录；否则，创建一个新的对话列表
     *
     * @param conversationId 对话ID，用于定位对话文件
     * @return 返回对话的消息列表，如果文件不存在则返回空列表
     */
    private List<Message> getOrCreateConversation(String conversationId) {
        // 获取对话文件对象
        File conversationFile = getConversationFile(conversationId);
        // 初始化消息列表
        List<Message> messages = new ArrayList<>();

        // 检查对话文件是否存在
        if (conversationFile.exists()) {
            try (Input input = new Input(new FileInputStream(conversationFile))) {
                // 从文件中读取对话记录到列表中
                messages = KRYO.readObject(input, ArrayList.class);
            } catch (IOException e) {
                // 打印异常信息，便于调试
                e.printStackTrace();
            }
        }
        // 返回消息列表
        return messages;
    }

    /**
     * 保存对话到文件
     *
     * @param conversationId 对话ID，用于确定对话文件
     * @param messages       消息列表，包含对话中的所有消息
     */
    private void saveConversation(String conversationId, List<Message> messages) {
        // 获取对话文件对象
        File conversationFile = getConversationFile(conversationId);

        // 创建文件输出流
        try (Output output = new Output(new FileOutputStream(conversationFile))) {
            // 将消息列表写入文件
            KRYO.writeObject(output, messages);
        } catch (IOException e) {
            // 打印异常信息，便于调试
            e.printStackTrace();
        }
    }
}
